// Copyright(c) 2016 Ruoyu Fan (Windy Darian), Xueyin Wan
// MIT License.

#include "Vulkan_util.h"

#include "Context.h"

#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>

uint32_t findMemoryType(uint32_t type_filter, vk::MemoryPropertyFlags properties, vk::PhysicalDevice physical_device)
{
	auto memory_properties = physical_device.getMemoryProperties();

	for (uint32_t i = 0; i < memory_properties.memoryTypeCount; i++)
	{
		bool type_supported = (type_filter & (1 << i)) != 0;
		bool properties_supported = ((memory_properties.memoryTypes[i].propertyFlags & properties) == properties);
		if (type_supported && properties_supported)
		{
			return i;
		}
	}

	throw std::runtime_error("Failed to find suitable memory type!");
}

VUtility::VUtility(const VContext & context)
	: context(&context)
	, physical_device(context.getPhysicalDevice())
	, device(context.getDevice())
	, graphics_queue(context.getGraphicsQueue())
	, compute_queue(context.getComputeQueue())
	, present_queue(context.getPresentQueue())
	, graphics_queue_command_pool(context.getGraphicsCommandPool())
	, compute_queue_command_pool(context.getComputeCommandPool())
{}

vk::SurfaceFormatKHR VUtility::chooseSwapSurfaceFormat(const std::vector<vk::SurfaceFormatKHR>& available_formats)
{
	// When free to choose format
	if (available_formats.size() == 1 && available_formats[0].format == vk::Format::eUndefined)
	{
		return{ vk::Format::eB8G8R8A8Unorm, vk::ColorSpaceKHR::eSrgbNonlinear};
	}

	for (const auto& available_format : available_formats)
	{
		// prefer 32bits RGBA color with SRGB support
		if (available_format.format == vk::Format::eB8G8R8A8Unorm && available_format.colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear)
		{
			return available_format;
		}
	}

	// TODO: Rank how good the formats are and choose the best?

	return available_formats[0];
}

vk::PresentModeKHR VUtility::chooseSwapPresentMode(const std::vector<vk::PresentModeKHR>& available_present_modes)
{
	for (const auto& available_present_mode : available_present_modes)
	{
		if (available_present_mode == vk::PresentModeKHR::eMailbox)
		{
			return available_present_mode;
		}
	}

	return vk::PresentModeKHR::eFifo;
}

vk::Extent2D VUtility::chooseSwapExtent(const vk::SurfaceCapabilitiesKHR& capabilities)
{
	// The swap extent is the resolution of the swap chain images
	if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max())
	{
		return capabilities.currentExtent;
	}
	else
	{
		int window_framebuffer_width, window_framebuffer_height;
		std::tie(window_framebuffer_width, window_framebuffer_height) = context->getWindowFrameBufferSize();
		vk::Extent2D actual_extent = { (uint32_t)window_framebuffer_width, (uint32_t)window_framebuffer_height };

		actual_extent.width = std::max(capabilities.minImageExtent.width
			, std::min(capabilities.maxImageExtent.width, actual_extent.width));
		actual_extent.height = std::max(capabilities.minImageExtent.height
			, std::min(capabilities.maxImageExtent.height, actual_extent.height));

		return actual_extent;
	}
}

vk::Format VUtility::findSupportedFormat(const std::vector<vk::Format>& candidates
	, vk::ImageTiling tiling, vk::FormatFeatureFlags features)
{
	for (vk::Format format : candidates)
	{
		auto props = physical_device.getFormatProperties(format);

		if (tiling == vk::ImageTiling::eLinear && (props.linearTilingFeatures & features) == features)
		{
			return format;
		}
		else if (tiling == vk::ImageTiling::eOptimal && (props.optimalTilingFeatures & features) == features)
		{
			return format;
		}
	}

	throw std::runtime_error("Failed to find supported format!");
}

std::tuple<VRaii<vk::Buffer>, VRaii<vk::DeviceMemory>> VUtility::createBuffer(vk::DeviceSize size, vk::BufferUsageFlags usage, vk::MemoryPropertyFlags property_bits
	, int sharing_queue_family_index_a, int sharing_queue_family_index_b)
{
	auto buffer_info = vk::BufferCreateInfo()
		.setSize(size)
		.setUsage(usage);

	std::array<uint32_t, 2> indices;
	if (sharing_queue_family_index_a >= 0 && sharing_queue_family_index_b >= 0 && sharing_queue_family_index_a != sharing_queue_family_index_b)
	{
		indices = { static_cast<uint32_t>(sharing_queue_family_index_a) , static_cast<uint32_t>(sharing_queue_family_index_b) };
		buffer_info.setSharingMode(vk::SharingMode::eConcurrent)
				.setQueueFamilyIndexCount(indices.size())
				.setPQueueFamilyIndices(indices.data());
	}
	else
	{
		buffer_info.setSharingMode(vk::SharingMode::eExclusive);
	}
	auto buffer = device.createBuffer(buffer_info);

	// allocate memory for buffer
	auto memory_req = device.getBufferMemoryRequirements(buffer);

	auto typeIndex = findMemoryType(memory_req.memoryTypeBits
			, property_bits
			, physical_device);
	auto memory_alloc_info = vk::MemoryAllocateInfo(memory_req.size,typeIndex);

	auto buffer_memory = device.allocateMemory(memory_alloc_info);

	device.bindBufferMemory(buffer,buffer_memory,0);

	auto raii_buffer_deleter = [this](vk::Buffer& obj)
	{
		device.destroyBuffer(obj);
	};

	auto raii_memory_deleter = [this](vk::DeviceMemory& obj)
	{
		device.freeMemory(obj);
	};

	return std::make_tuple(VRaii<vk::Buffer>(buffer, raii_buffer_deleter)
		, VRaii<vk::DeviceMemory>(buffer_memory, raii_memory_deleter));
}

void VUtility::copyBuffer(vk::Buffer src_buffer, vk::Buffer dst_buffer, vk::DeviceSize size, vk::DeviceSize src_offset, vk::DeviceSize dst_offset)
{
	vk::CommandBuffer copy_command_buffer = beginSingleTimeCommands();

	recordCopyBuffer(copy_command_buffer, src_buffer, dst_buffer, size, src_offset, dst_offset);

	endSingleTimeCommands(copy_command_buffer);
}

std::tuple<VRaii<vk::Image>, VRaii<vk::DeviceMemory>> VUtility::createImage(uint32_t image_width, uint32_t image_height
	, vk::Format format, vk::ImageTiling tiling
	, vk::ImageUsageFlags usage, vk::MemoryPropertyFlags memory_properties)
{
    auto image_info = vk::ImageCreateInfo()
        .setImageType(vk::ImageType::e2D)
        .setExtent({image_width,image_height,1})
        .setMipLevels(1)
        .setArrayLayers(1)
        .setFormat(format)
        .setTiling(tiling)
        .setInitialLayout(vk::ImageLayout::ePreinitialized)
        .setUsage(usage)
        .setSharingMode(vk::SharingMode::eExclusive);

    auto image = device.createImage(image_info);

	// allocate image memory
    auto memory_req = device.getImageMemoryRequirements(image);
    auto typeIndex = findMemoryType(memory_req.memoryTypeBits
            , memory_properties
            , physical_device);

    auto alloc_info = vk::MemoryAllocateInfo(memory_req.size,typeIndex);

    auto memory = device.allocateMemory(alloc_info);
    device.bindImageMemory(image,memory,0);

	auto raii_image_deleter = [this](vk::Image& obj)
	{
		device.destroyImage(obj);
	};

	auto raii_memory_deleter = [this](vk::DeviceMemory& obj)
	{
		device.freeMemory(obj);
	};

	return std::make_tuple(VRaii<vk::Image>(image, raii_image_deleter), VRaii<vk::DeviceMemory>(memory, raii_memory_deleter));
}

void VUtility::copyImage(vk::Image src_image, vk::Image dst_image, uint32_t width, uint32_t height)
{
	vk::CommandBuffer command_buffer = beginSingleTimeCommands();

	recordCopyImage(command_buffer, src_image, dst_image, width, height);

	endSingleTimeCommands(command_buffer);
}

void VUtility::transitImageLayout(vk::Image image, vk::ImageLayout old_layout, vk::ImageLayout new_layout)
{
	vk::CommandBuffer command_buffer = beginSingleTimeCommands();

	recordTransitImageLayout(command_buffer, image, old_layout, new_layout);

	endSingleTimeCommands(command_buffer);

}

VRaii<vk::ImageView> VUtility::createImageView(vk::Image image, vk::Format format, vk::ImageAspectFlags aspect_mask)
{
	auto viewInfo = vk::ImageViewCreateInfo()
		.setImage(image)
		.setViewType(vk::ImageViewType::e2D)
		.setFormat(format)
		.setSubresourceRange({aspect_mask,0,1,0,1});

	auto deleter = [this](vk::ImageView& obj) {device.destroyImageView(obj); };
	auto view = device.createImageView(viewInfo);
	return VRaii<vk::ImageView>(view,deleter);
}

std::tuple<VRaii<vk::Image>, VRaii<vk::DeviceMemory>, VRaii<vk::ImageView>> VUtility::loadImageFromFile(std::string path)
{
	// TODO: maybe move to vulkan_util or a VulkanDevice class

	// load image file
	int tex_width, tex_height, tex_channels;

	stbi_uc * pixels = stbi_load(path.c_str()
		, &tex_width, &tex_height
		, &tex_channels
		, STBI_rgb_alpha);

	vk::DeviceSize image_size = tex_width * tex_height * 4;
	if (!pixels)
	{
		throw std::runtime_error("Failed to load image" + path);
	}

	// create staging image memory
	VRaii<vk::Image> staging_image;
	VRaii<vk::DeviceMemory> staging_image_memory;
	std::tie(staging_image, staging_image_memory) = createImage(
		tex_width, tex_height
		, vk::Format::eR8G8B8A8Unorm
		, vk::ImageTiling::eLinear
		, vk::ImageUsageFlagBits::eTransferSrc
		, vk::MemoryPropertyFlags()|vk::MemoryPropertyFlagBits::eHostVisible|vk::MemoryPropertyFlagBits::eHostCoherent
	);

	// copy image to staging memory
	void* data;
	vkMapMemory(device, staging_image_memory.get(), 0, image_size, 0, &data);
	memcpy(data, pixels, (size_t)image_size);
	vkUnmapMemory(device, staging_image_memory.get());

	// free image in memory
	stbi_image_free(pixels);

	VRaii<vk::Image> image;
	VRaii<vk::DeviceMemory> image_memory;
	// create texture image
	std::tie(image, image_memory) = createImage(
		tex_width, tex_height
		, vk::Format::eR8G8B8A8Unorm
		, vk::ImageTiling::eOptimal
		, vk::ImageUsageFlags()|vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eSampled
		, vk::MemoryPropertyFlagBits::eDeviceLocal
	);

	// TODO: doing the steps asynchronously by using a single command buffer
	auto command_buffer = beginSingleTimeCommands();

	recordTransitImageLayout(command_buffer, staging_image.get(), vk::ImageLayout::ePreinitialized, vk::ImageLayout::eTransferSrcOptimal);
	recordTransitImageLayout(command_buffer, image.get(), vk::ImageLayout::ePreinitialized, vk::ImageLayout::eTransferDstOptimal);
	recordCopyImage(command_buffer, staging_image.get(), image.get(), tex_width, tex_height);
	recordTransitImageLayout(command_buffer, image.get(), vk::ImageLayout::eTransferDstOptimal, vk::ImageLayout::eShaderReadOnlyOptimal);

	endSingleTimeCommands(command_buffer);


	// Create image view
	auto image_view = createImageView(image.get(), vk::Format::eR8G8B8A8Unorm, vk::ImageAspectFlagBits::eColor);

	return std::make_tuple(std::move(image), std::move(image_memory), std::move(image_view));
}

VRaii<std::vector<vk::CommandBuffer>> VUtility::createCommandBuffers(uint32_t count,vk::CommandBufferLevel level)
{
	auto deleter = [this](std::vector<vk::CommandBuffer>& obj){
		device.freeCommandBuffers(graphics_queue_command_pool,obj);
	};
	auto alloc_info = vk::CommandBufferAllocateInfo(graphics_queue_command_pool,level,count);
	auto command_buffer = device.allocateCommandBuffers(alloc_info);
	return {command_buffer,deleter};
}

// create a temperorary command buffer for one-time use
// and begin recording
vk::CommandBuffer VUtility::beginSingleTimeCommands()
{
    auto alloc_info = vk::CommandBufferAllocateInfo(graphics_queue_command_pool,vk::CommandBufferLevel::ePrimary,1);
    auto command_buffer = device.allocateCommandBuffers(alloc_info);
    auto begin_info = vk::CommandBufferBeginInfo(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);
    command_buffer[0].begin(begin_info);
	return command_buffer[0];
}

// End recording the single time command, submit then wait for execution and destroy the buffer
void VUtility::endSingleTimeCommands(vk::CommandBuffer command_buffer)
{
    command_buffer.end();

    auto submit_info = vk::SubmitInfo()
        .setCommandBufferCount(1)
        .setPCommandBuffers(&command_buffer);

    graphics_queue.submit(submit_info,VK_NULL_HANDLE);
    graphics_queue.waitIdle();

    device.freeCommandBuffers(graphics_queue_command_pool,1,&command_buffer);
}

void VUtility::recordCopyBuffer(vk::CommandBuffer command_buffer, vk::Buffer src_buffer, vk::Buffer dst_buffer, vk::DeviceSize size, vk::DeviceSize src_offset, vk::DeviceSize dst_offset)
{
    auto copy_region = vk::BufferCopy(src_offset,dst_offset,size);
    command_buffer.copyBuffer(src_buffer,dst_buffer,copy_region);
}

void VUtility::recordCopyImage(vk::CommandBuffer command_buffer, vk::Image src_image, vk::Image dst_image, uint32_t width, uint32_t height)
{
    auto sub_besource = vk::ImageSubresourceLayers(vk::ImageAspectFlagBits::eColor,0,0,1);

    auto region = vk::ImageCopy(sub_besource,{0,0,0},sub_besource,{0,0,0},{width,height,1});

    command_buffer.copyImage(src_image,vk::ImageLayout::eTransferSrcOptimal,
        dst_image,vk::ImageLayout::eTransferDstOptimal,
        region);
}

void VUtility::recordTransitImageLayout(vk::CommandBuffer command_buffer, vk::Image image, vk::ImageLayout old_layout, vk::ImageLayout new_layout)
{
	// barrier is used to ensure a buffer has finished writing before
	// reading as weel as doing transition
    auto barrier = vk::ImageMemoryBarrier()
        .setOldLayout(old_layout)
        .setNewLayout(new_layout)
        .setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)//TODO: when converting the depth attachment for depth pre pass this is not correct
        .setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
        .setImage(image);

	if (new_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal || old_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal)
	{
		barrier.subresourceRange.aspectMask = vk::ImageAspectFlagBits::eDepth;
	}
	else
	{
		barrier.subresourceRange.aspectMask = vk::ImageAspectFlagBits::eColor;
	}

	barrier.subresourceRange.baseMipLevel = 0;
	barrier.subresourceRange.levelCount = 1;
	barrier.subresourceRange.baseArrayLayer = 0;
	barrier.subresourceRange.layerCount = 1;

	if (old_layout == vk::ImageLayout::ePreinitialized && new_layout == vk::ImageLayout::eTransferSrcOptimal)
	{
		// dst must wait on src
		barrier.setSrcAccessMask(vk::AccessFlagBits::eHostWrite);
		barrier.setDstAccessMask(vk::AccessFlagBits::eTransferRead);
	}
	else if (old_layout == vk::ImageLayout::ePreinitialized && new_layout == vk::ImageLayout::eTransferDstOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlagBits::eHostWrite);
		barrier.setDstAccessMask(vk::AccessFlagBits::eTransferWrite);
	}
	else if (old_layout == vk::ImageLayout::eTransferDstOptimal && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlagBits::eTransferWrite);
		barrier.setDstAccessMask(vk::AccessFlagBits::eShaderRead);
	}
	else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlags());
		barrier.setDstAccessMask(vk::AccessFlags()|vk::AccessFlagBits::eDepthStencilAttachmentRead
			| vk::AccessFlagBits::eDepthStencilAttachmentWrite);
	}
	else if (old_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlags()|vk::AccessFlagBits::eDepthStencilAttachmentRead
			| vk::AccessFlagBits::eDepthStencilAttachmentWrite);
		barrier.setDstAccessMask(vk::AccessFlagBits::eShaderRead);
	}
	else if (old_layout == vk::ImageLayout::eShaderReadOnlyOptimal && new_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlagBits::eShaderRead);
		barrier.setDstAccessMask(vk::AccessFlags()|vk::AccessFlagBits::eDepthStencilAttachmentRead
			| vk::AccessFlagBits::eDepthStencilAttachmentWrite);
	}
	else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal)
	{
		barrier.setSrcAccessMask(vk::AccessFlags());
		barrier.setDstAccessMask(vk::AccessFlagBits::eShaderRead);
	}
	else
	{
		throw std::invalid_argument("unsupported layout transition!");
	}

    command_buffer.pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe,
                                   vk::PipelineStageFlagBits::eTopOfPipe,
                                   vk::DependencyFlags(),
                                   nullptr, nullptr, barrier);
}
